import ttkbootstrap as tk

from tkinter import font as tkfont
from ttkbootstrap.constants import *
from ttkbootstrap.utility import enable_high_dpi_awareness
from PIL import Image, ImageTk
from sys import exit

from src.common import get_screenshoot, cut_image


class MaskWindow:
    def __init__(self, parent, filepath):
        self.close_callback = None
        self.save_path = filepath

        self.root = tk.Toplevel(parent)
        self.key = None
        hex_to_rgba = lambda s: tuple(int(s.lstrip('#')[i:i + 2], 16) for i in range(0, 8, 2))

        self.mask_bg = hex_to_rgba('#1d2129e2')
        self.root.overrideredirect(True)

        # 设置30秒后自动取消
        self.root.after(30000, self.end)
        
        # 获取屏幕尺寸
        self.screen_width = parent.winfo_screenwidth()
        self.screen_height = parent.winfo_screenheight()
        
        # 设置窗口全屏和透明
        self.root.geometry(f'{self.screen_width}x{self.screen_height}+0+0')
        self.root.attributes('-alpha', 1.0)
        self.root.attributes('-topmost', True)

        self.font = tkfont.Font(family='阿里妈妈方圆体 VF Medium', size=15)
        
        # 创建全屏Canvas
        self.canvas = tk.Canvas(
            self.root,
            width=self.screen_width,
            height=self.screen_height,
            highlightthickness=0,
            cursor="crosshair"
        )
        self.canvas.pack(fill='both', expand=True)
        
        # 创建背景图像
        self.screenshoot = get_screenshoot()
        self.basemaop = self.add_mask(self.screenshoot, (0,0,0,0), self.mask_bg)
        self.canvas.create_image(0, 0, image=self.basemaop, anchor='nw')
        self.draw_hint(0, 0, 0, 0)

        # 截图区域
        self.start_x = None
        self.start_y = None
        self.current_rect = None
        self.selection_area = None

        # 绑定事件
        self.canvas.bind('<Button-1>', self.on_press)
        self.canvas.bind('<B1-Motion>', self.on_drag)
        self.canvas.bind('<ButtonRelease-1>', self.on_release)
        self.canvas.bind('<Button-3>', lambda e: self.end())
        self.root.bind('<Key>', self.on_key)

    def create_rounded_rectangle(self, x1, y1, x2, y2, radius, **kwargs):
        # 创建圆角矩形的点
        points = [
            x1 + radius, y1,
            x2 - radius, y1,
            x2, y1,
            x2, y1 + radius,
            x2, y2 - radius,
            x2, y2,
            x2 - radius, y2,
            x1 + radius, y2,
            x1, y2,
            x1, y2 - radius,
            x1, y1 + radius,
            x1, y1
        ]
        return self.canvas.create_polygon(points, smooth=True, **kwargs)

    def on_press(self, event):
        self.start_x = event.x
        self.start_y = event.y

        # 清除之前的选择
        if self.current_rect:
            self.canvas.delete(self.current_rect)
        if self.selection_area:
            self.canvas.delete(self.selection_area)

        # 创建选区
        self.selection_area = self.canvas.create_rectangle(
            self.start_x, self.start_y, self.start_x, self.start_y,
            fill='', outline=''
        )

        # 创建选区边框
        self.current_rect = self.canvas.create_rectangle(
            self.start_x, self.start_y, self.start_x, self.start_y,
            outline='#1890ff',
            width=2,
            fill=''
        )

    def on_drag(self, event):
        if self.start_x is None or self.start_y is None:
            return
        if self.current_rect and self.selection_area:
            # 更新选区大小
            x1 = min(self.start_x, event.x)
            y1 = min(self.start_y, event.y)
            x2 = max(self.start_x, event.x)
            y2 = max(self.start_y, event.y)

            # 更新遮罩和选区
            self.canvas.delete('all')  # 清除所有内容

            # 重新绘制背景图
            self.basemaop = self.add_mask(self.screenshoot, (x1, y1, x2, y2), self.mask_bg)
            self.canvas.create_image(0, 0, image=self.basemaop, anchor='nw')

            # 绘制选区边框（带圆角）
            self.current_rect = self.create_rounded_rectangle(
                x1, y1, x2, y2,
                radius=8,  # 8px圆角
                outline='#6AA1FF',
                width=4,
                fill='',
                dash=(4, 4)
            )

            # 绘制提示文字
            self.draw_hint(x1, y1, x2, y2)

    def draw_hint(self, x1, y1, x2, y2):
        if x1 > 384 or y2 < (self.screen_height - 176):
            # 创建半透明黑色背景，不使用stipple
            text_bg = self.canvas.create_rectangle(
                24, self.screen_height - 176,
                384, self.screen_height - 16,
                fill='#1d2129',  # 使用相同的遮罩颜色
                outline=''  # 无边框
            )

            self.canvas.create_text(
                48, self.screen_height - 136,
                text="按住鼠标左键并拖动鼠标",
                fill='white',
                font=self.font,
                anchor='w'
            )
            self.canvas.create_text(
                148, self.screen_height - 98,
                text="来选择截屏区域",
                fill='white',
                font=self.font,
                anchor='w'
            )
            self.canvas.create_text(
                48, self.screen_height - 56,
                text="单机鼠标右键直接退出",
                fill='white',
                font=self.font,
                anchor='w'
            )

    def add_mask(self, image, box, mask_bg):
        """
        为截图添加遮罩效果
        :param image: 原始PhotoImage对象
        :param box: 选择区域 (x1, y1, x2, y2)
        :param mask_bg: 遮罩颜色 (r, g, b, a)
        :return: 处理后的PhotoImage对象
        """

        # 确保图像是RGBA模式
        if image.mode != 'RGBA':
            image = image.convert('RGBA')

        # 创建遮罩层
        mask = Image.new('RGBA', image.size, mask_bg)

        # 如果有选择区域，则清除该区域的遮罩
        if box != (0, 0, 0, 0):
            x1, y1, x2, y2 = box
            # 创建透明区域
            transparent = Image.new('RGBA', (x2 - x1, y2 - y1))
            mask.paste(transparent, (x1, y1))

        # 将原图和遮罩层合并
        result = Image.alpha_composite(image, mask)

        # 调整图像大小以匹配窗口大小
        # result = result.resize((self.screen_width, self.screen_height))

        # 转为RGB
        result = result.convert('RGB')

        # 转换回PhotoImage
        return ImageTk.PhotoImage(result)

    def on_release(self, event):
        if self.start_x is None or self.start_y is None:
            return
            
        # 确保有最小选择区域
        if abs(event.x - self.start_x) < 5 or abs(event.y - self.start_y) < 5:
            return
            
        # 获取选择区域坐标
        x1 = min(self.start_x, event.x)
        y1 = min(self.start_y, event.y)
        x2 = max(self.start_x, event.x)
        y2 = max(self.start_y, event.y)
        
        # 截图
        screenshot = cut_image(box=(x1, y1, x2, y2))
        
        # 保存截图
        screenshot.save(self.save_path)
        
        self.end()
        
    def on_key(self, event):
        # 处理键盘事件
        if event.keysym == 'Escape':
            self.end()

    def show(self):
        self.root.mainloop()

    def end(self):
        if self.close_callback:
            self.close_callback()

    def setCloseCallback(self, callback):
        self.close_callback = callback
    
    def close(self):
        self.root.destroy()


class ScreenshotDialog:
    
    def __init__(self, parent):
        self.__parent = parent
    
    def initUI(self, params):
        try:
            self.save_path = params.filepath
            # 获取屏幕信息并且在每一个屏幕上创建一个窗口
            self.windows = []
            self.bottom_images = []
            # for screen in self.screens:
            window = MaskWindow(self.__parent, self.save_path)
            window.setCloseCallback(self.close)
            self.windows.append(window)

        except Exception as e:
            print(f"初始化截屏错误: {e}")
            import traceback
            traceback.print_exc()

    def show(self):
        for window in self.windows:
            window.show()
    
    def close(self):
        for window in self.windows:
            window.close()
        exit(0)


def create(parser):
    parser.add_argument("--filepath", "-f", type=str, default='', help="截图保存路径")
    params = parser.parse_args()

    enable_high_dpi_awareness()
    root = tk.Window(alpha=0.0, overrideredirect=True, maxsize=(1, 1))
    root.place_window_center()
    screenshot = ScreenshotDialog(root)
    screenshot.initUI(params)
    screenshot.show()
